import cv2
import os
import math

def extract_frames_uniformly(video_path, n_frames, output_dir="extracted_frames"):
    """
    从视频中均匀提取n帧
    
    参数:
    video_path: 视频文件路径
    n_frames: 要提取的帧数
    output_dir: 输出目录，默认为"extracted_frames"
    """
    
    # 检查视频文件是否存在
    if not os.path.exists(video_path):
        print(f"错误：视频文件 {video_path} 不存在")
        return False
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建输出目录: {output_dir}")
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    
    if not cap.isOpened():
        print(f"错误：无法打开视频文件 {video_path}")
        return False
    
    # 获取视频信息
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    duration = total_frames / fps
    
    print(f"视频信息:")
    print(f"  总帧数: {total_frames}")
    print(f"  帧率: {fps:.2f} FPS")
    print(f"  时长: {duration:.2f} 秒")
    print(f"  要提取的帧数: {n_frames} (排除开头和结尾帧)")
    
    # 如果要提取的帧数大于可用帧数，进行调整
    max_available = total_frames - 2  # 排除开头和结尾后的最大可用帧数
    if n_frames > max_available:
        print(f"警告：要提取的帧数({n_frames})大于可用帧数({max_available})，调整为{max_available}")
        n_frames = max_available
    
    # 计算要提取的帧的位置（排除开头和结尾帧）
    # 定义可用帧范围：跳过第一帧和最后一帧
    start_frame = 1  # 跳过第0帧（开头）
    end_frame = total_frames - 2  # 跳过最后一帧（结尾）
    available_frames = end_frame - start_frame + 1
    
    print(f"可用帧范围: 第{start_frame}帧 到 第{end_frame}帧 (共{available_frames}帧)")
    
    # 检查可用帧数是否足够
    if available_frames < n_frames:
        print(f"警告：可用帧数({available_frames})少于要提取的帧数({n_frames})，调整为{available_frames}")
        n_frames = available_frames
    
    if n_frames == 1:
        # 如果只取1帧，取中间位置的帧
        frame_indices = [(start_frame + end_frame) // 2]
    else:
        # 在可用帧范围内均匀分布
        frame_indices = []
        for i in range(n_frames):
            # 在start_frame到end_frame之间均匀分布
            frame_idx = start_frame + int(i * (end_frame - start_frame) / (n_frames - 1))
            frame_indices.append(frame_idx)
    
    print(f"将提取以下帧(排除开头和结尾): {frame_indices}")
    
    # 提取帧
    extracted_count = 0
    for i, frame_idx in enumerate(frame_indices):
        # 设置视频位置到指定帧
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
        
        # 读取帧
        ret, frame = cap.read()
        if ret:
            # 计算时间戳
            timestamp = frame_idx / fps
            
            # 生成文件名
            filename = f"frame_{i+1:03d}_at_{timestamp:.2f}s.jpg"
            output_path = os.path.join(output_dir, filename)
            
            # 保存帧
            success = cv2.imwrite(output_path, frame)
            if success:
                print(f"保存帧 {i+1}/{n_frames}: {filename}")
                extracted_count += 1
            else:
                print(f"错误：无法保存帧到 {output_path}")
        else:
            print(f"警告：无法读取第 {frame_idx} 帧")
    
    # 释放资源
    cap.release()
    
    print(f"\n提取完成！成功提取了 {extracted_count} 帧到目录: {output_dir}")
    return True

# 配置参数 - 在这里修改你的设置
if __name__ == "__main__":
    # ===== 配置区域 - 请修改以下参数 =====
    VIDEO_PATH = "/home/shuai/Videos/避障特写_第一视角.mp4"  # 替换为你的视频文件路径
    N_FRAMES = 10  # 要提取的帧数
    OUTPUT_DIR = "/home/shuai/Videos/extracted_frames1"  # 输出目录，可选
    
    # ===== 执行提取 =====
    print("开始提取视频帧...")
    success = extract_frames_uniformly(VIDEO_PATH, N_FRAMES, OUTPUT_DIR)
    
    if success:
        print("任务完成！")
    else:
        print("任务失败，请检查错误信息。")
